# -*- coding: utf-8 -*-
"""
Wcrypto Module 
Function used for crypto
implemented with PyCrypto
Make you don't care about Key Size,Block Size
hash your key into the key size and format your data into block size
set Crypto method like this 'AES-CBC-32'
AES is a Cipher
CBC is a Cipher Mode
32 is Key Size
default sets 'AES','CBC',32
Here is a list of Ciphers and Modes you can used
-----------Ciphers-----------------------
Cipher          Key Size/Block Size
AES             16, 24, or 32 bytes/16 bytes
ARC2            Variable/8 bytes
ARC4            Variable/1 bytes
Blowfish        Variable/8 bytes
CAST            5-16 bytes/8 bytes
DES             8 bytes/8 bytes
DES3 (Triple DES)   16 bytes/8 bytes
IDEA            16 bytes/8 bytes###removed by GAE
RC5             at least 2 bytes/8 bytes###removed by GAE
XOR             1-32bytes/1 bytes
-----------------------------------------
New PyCrypto version 2.1.0alpha1 has removed the IDEA and RC5 modules due to software patents.(I also support)
-----------Modes-------------------------
ECB fastest     0.391s  may be dangerous(unsafe)
CBC middle      0.453s  default and suggest
CFB slowest     5.375s
CTR    cann't used      #undocument by PyCrypto
OFB             0.437s  #undocument by PyCrypto
PGP             0.438s  #removed by GAE
-------------------------------------------

1.the cipher which block size is 1 DO NOT have modes
2.RC5 with keysize=1,17 will cause a MemmoryError
3.if you use those Ciphers and Modes which are removed by GAE on GAE will  
   raise a ImportError
Example:
akey=Wcrypto.Wcrypto('AES-CBC-32') #for init 
encrypted_data=akey.encrypt(data,key) #encrypt
dosomethingwith(encrypted_data) #send,store or something else
data=akey.decrypt(encrypted_data,key) #decrypt
"""

__author__ = 'redsxz(bronze1man@gmail.com)'
__version__ = '0.3.9'
__all__ = ["Wcrypto","RandomMethod"]

import hashlib
import os
#import Crypto.Cipher.AES
#import Crypto.Cipher.ARC2
#import Crypto.Cipher.ARC4
#import Crypto.Cipher.Blowfish
#import Crypto.Cipher.CAST
#import Crypto.Cipher.DES
#import Crypto.Cipher.DES3
#import Crypto.Cipher.IDEA
#import Crypto.Cipher.RC5
#import Crypto.Cipher.XOR

class Wcrypto:
  
  cipher='AES'
  mode='CBC'
  BlockSize=16
  keysize=0
  ishashkey=True
  _usepurexor=False
  cipherlib=None
  ciphermodeid=None
  _new=None
  
  
  
  def __init__(self,smode='AES-CBC-32'):
    self.initstr(smode)
  
  def init(self,cipher='AES',mode='CBC',keysize=32,ishashkey=True):
    """
cipher:Crypto method
mode:the mode of Crypto method
keysize:if ishashkey=True it means hashkeysize
        if ishashkey=False it means nothing
ishashkey:whether hash the imput key
"""
    self.cipher=cipher;
    self.mode=mode;
    self.keysize=keysize;
    self.ishashkey=ishashkey
    
    #Invalid cipher
    if not(cipher in self._ListofCiphers):
      raise ValueError("invalid cipher:"+cipher)
    self.BlockSize=self._ListofBlockSize[cipher]
    #invalid mode
    if not(mode in self._ListofModes):self.mode='CBC'
    if (self.BlockSize==1):self.mode=''
    if (ishashkey):
      if not(cipher in self._ListofKeySize):
      #keysize is not need to fixed
      #invalid keysize out of range
        if (keysize<=0):
          self.keysize=32
      #invalid keysize Not in dict
      else:#keysize must be fixed
        if not(keysize in self._ListofKeySize[cipher]):
          self.keysize = self._ListofKeySize[cipher][-1]
    else:#ishashkey is False
      pass#do nothing
    #follow will cause a Memmory Error
    if (self.cipher=='RC5')and(self.keysize==1 or self.keysize==17):
        self.keysize=32
    
    
    #try to import Crypto.Cipher.xxxx
    try:
      #maybe run any code
      exec 'import Crypto.Cipher.'+cipher
      self._usepurexor=False
      self.cipherlib=getattr(Crypto.Cipher,cipher)
      if (self.BlockSize!=1):
          self.ciphermodeid=getattr(self.cipherlib,'MODE_'+self.mode)
      self._new=self.cipherlib.new
    except ImportError:
      if (cipher=='XOR'):self._usepurexor=True
      else:raise ImportError('Crypto.Cipher.'+cipher+'no found')
  
  
  def initstr(self,smode='AES-CBC-32'):
    if (smode==''):return self.init()
    slist=smode.split('-')
    lenl=len(slist)
    cipher='';mode='CBC';keysize=32
    if (lenl>=1):cipher=slist[0]
    if (lenl>=2):mode=slist[1]
    try:
      if (lenl>=3):keysize=int(slist[2])
    except ValueError:
      raise TypeError(smode+"[the last part must be a int]")
    if (keysize>0):return self.init(cipher,mode,keysize,True)
    else:return self.init(cipher,mode,keysize,False)
  
  
  def encrypt(self,data,key):
    khash=hashlib.sha512(key).digest()
    if (self.ishashkey):_key=khash[:self.keysize]
    else:_key=key
    if (self._usepurexor):return self._xor(data,key);
    #insert data to block size
    if (self.BlockSize!=1):
      #block Crypto
      a=len(data)+1
      add=((self.BlockSize-a%self.BlockSize)%self.BlockSize)
      edata=data + os.urandom(add) + chr(add)
    #new a key
      tkey=self._new(_key,self.ciphermodeid,
                khash[-self.BlockSize:])
      return tkey.encrypt(edata)
    else:#stream Crypto
      if (self.cipher=='XOR')and(_key==''): return data
      return self._new(_key).encrypt(data)
  
  
  def decrypt(self,data,key):
    khash=hashlib.sha512(key).digest()
    if (self.ishashkey):_key=khash[:self.keysize]
    else:_key=key
    if (self._usepurexor):return self._xor(data,key);
    #insert data to block size
    if (self.BlockSize!=1):
      #new a key
      tkey=self._new(_key,self.ciphermodeid,
                khash[-self.BlockSize:])
      edata=tkey.decrypt(data)
      a=ord(edata[-1])
        #Wrong key,return a Wrong data
      if (a>self.BlockSize):return edata
      return edata[:-(a+1)]
    else:#stream Crypto
      if (self.cipher=='XOR')and(_key==''): return data
      return self._new(_key).decrypt(data)
  
  def realmodestr(self):
    return self.cipher+'-'+self.mode+'-'+str(self.keysize)
  
  #in case No Crypto ,wo can also XOR with pure Python
  def _xor(self,data,key):
    import itertools
    if (key==''): return data
    xorsize=1000
    start=0;end=len(data)
    ss=['']*(end/xorsize+2);iss=0
    key = [ord(i) for i in key]
    ck=itertools.cycle(key)
    while start<end:
        dend=start+xorsize
        if (dend>end):dend=end
        dd = [ord(d)^k for d,k in itertools.izip(data[start:dend], ck)]
        ss[iss]=''.join(map(chr, dd));iss+=1
        start=dend
    return ''.join(ss)
  
  
  _ListofCiphers = ['AES','ARC2','ARC4','Blowfish','CAST',
                    'DES','DES3','IDEA','RC5','XOR']
  _ListofModes   = ['ECB','CBC','CFB','OFB','PGP']#Don't know how use CTR
  _ListofKeySize = {'AES':[16,24,32],
                    'DES':[8],
                    'DES3':[16],
                    'IDEA':[16],
                    'CAST':[5,6,7,8,9,10,11,12,13,14,15,16],
                    'XOR':range(1,33)}
  _ListofBlockSize={'AES':16,
                    'ARC2':8,
                    'ARC4':1,
                    'Blowfish':8,
                    'CAST':8,
                    'DES':8,
                    'DES3':8,
                    'IDEA':8,
                    'RC5':8,
                    'XOR':1}
  _ListofSuggestCipers=['AES','ARC2','ARC4','Blowfish','CAST',
                    'DES','DES3']
  _ListofSuggestModes=['ECB','CBC','CFB','OFB']
  
def RandomMethod(otherCiphers=False,otherMode=False,randomkeysize=False):
  import random
  #random Cipher
  if (otherCiphers==False):
    sc=random.choice(Wcrypto._ListofSuggestCipers)
  else: sc=random.choice(Wcrypto._ListofCipers)
  #random Mode
  if (Wcrypto._ListofBlockSize[sc]!=1):
    if (otherMode==False):sm=random.choice(Wcrypto._ListofSuggestModes)
    else: sm=random.choice(Wcrypto._ListofModes)
  else: sm=''
  #fixed keysize  
  if (sc in Wcrypto._ListofKeySize):
    keysize=Wcrypto._ListofKeySize[sc][-1]
  else:keysize=32
  return sc+'-'+sm+'-'+str(keysize)
  
